//
// Created by zhangcc on 2020/7/4.
//

#ifndef CRNNnet
#define CRNNnet

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <opencv2/opencv.hpp>

#include "base_crnn.h"

#include "../utils/cudaWrapper.h"
#include "../utils/ioHelper.h"

#define MAX_WORKSPACE_SIZE ((1ULL << 30) * 5) // 1 GB
#define MAX_BINDING_NUMS   (2)
#define DEBUG

class CrnnNet : public BaseCrnn {
public:
    CrnnNet(int device_id, bool trt_serialize = false):BaseCrnn(device_id, trt_serialize){};

    virtual int init(const std::string &model_file);

    std::vector <std::string> extract(const cv::Mat &inputs);

    ~CrnnNet();

private:
    int serialization(const std::string &model_file, const std::string &engine_file);

    std::unique_ptr <nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> mEngine{nullptr};

    std::unique_ptr <nvinfer1::IExecutionContext, nvinfer1::Destroy<nvinfer1::IExecutionContext>> context{nullptr};

    bool Input(cv::Mat &src, float *data);

    void *bindings[MAX_BINDING_NUMS]{0};
    cudaStream_t stream_t;
    cv::Mat input;
//    int device_id = 0;
//    bool trt_serialize;
};

#endif //CRNNnet
